package udpserver

import (
	"context"
	"fmt"
	"net"
	"nprobe/core"
	"nprobe/core/agent"
	"nprobe/core/aggregate"
	"syscall"
)

// UDPServerAgent icmp echo
type UDPServerAgent struct {
	agent.BaseAgent
	UDPConns []*net.UDPConn
}

func NewUDPServerAgent(name string, task *core.ProbeTask) (*UDPServerAgent, error) {

	udpConns, err := listenUDPs(task)
	if err != nil {
		fmt.Println("nprobe listen udp error: ", err)
		return nil, err
	}

	salts := core.NewSalts(task.MsgLen)

	ctx, cancelFunc := context.WithCancel(context.Background())
	receiverGroup := NewUDPServerReceiverGroup(ctx, name, task, 10000, salts)

	aggregator := aggregate.NewRingBufAggregator(ctx, name, task, nil, receiverGroup)

	return &UDPServerAgent{
		BaseAgent: agent.BaseAgent{
			Name:          name,
			Task:          task,
			SenderGroup:   nil,
			ReceiverGroup: receiverGroup,
			Aggregator:    aggregator,

			Ctx:        ctx,
			CancelFunc: cancelFunc,
		},
		UDPConns: udpConns,
	}, nil
}

func listenUDPs(task *core.ProbeTask) ([]*net.UDPConn, error) {
	var udpConns []*net.UDPConn
	for port := task.DstPortMin; port <= task.DstPortMax; port++ {
		conn, err := net.ListenUDP("udp4", &net.UDPAddr{Port: int(port)})
		if err != nil {
			for _, udpConn := range udpConns {
				udpConn.Close()
			}
			return nil, err
		}

		rawConn, err := conn.SyscallConn()
		if err != nil {
			for _, udpConn := range udpConns {
				udpConn.Close()
			}
			return nil, err
		}

		err = rawConn.Control(func(fd uintptr) {
			// 设置SO_REUSEADDR选项
			syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1)
		})

		udpConns = append(udpConns, conn)
	}

	return udpConns, nil
}

func (u *UDPServerAgent) StopAgent() error {
	if u != nil {
		u.BaseAgent.StopAgent()

		for _, conn := range u.UDPConns {
			conn.Close()
		}
	}
	return nil
}
